package crawler;

import org.jsoup.Jsoup;
import org.jsoup.nodes.Document;
import org.jsoup.nodes.Element;
import org.jsoup.select.Elements;
import utils.FileStorageUtil;

import java.io.IOException;
import java.util.Set;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

public class WebCrawler {
    private final Set<String> visitedUrls = ConcurrentHashMap.newKeySet();
    private final AtomicInteger englishDocCount = new AtomicInteger(0);
    private final AtomicInteger chineseDocCount = new AtomicInteger(0);
    private final int maxDocsPerLanguage = 500;
    private final ExecutorService executorService = Executors.newFixedThreadPool(10);
    private final BlockingQueue<String> urlQueue = new LinkedBlockingQueue<>();
    private static final String BASE_DATA_DIR = "data";
    private static final String RAW_DIR = "raw";
    private static final String PROCESSED_DIR = "processed";

    public void startCrawling() throws IOException, InterruptedException {
        // 初始化队列
        addSeedUrls();

        // 启动爬虫线程
        for (int i = 0; i < 10; i++) {
            executorService.execute(this::crawlWorker);
        }

        // 监控进度
        while (englishDocCount.get() < maxDocsPerLanguage ||
                chineseDocCount.get() < maxDocsPerLanguage) {
            System.out.printf("Progress: English %d/%d, Chinese %d/%d%n",
                    englishDocCount.get(), maxDocsPerLanguage,
                    chineseDocCount.get(), maxDocsPerLanguage);
            Thread.sleep(5000);
        }

        executorService.shutdown();
    }

    private void addSeedUrls() {
        // 英文种子URL
        urlQueue.add("http://en.people.cn/90777");
        urlQueue.add("http://en.people.cn/90780");
        urlQueue.add("https://www.chinadaily.com.cn/world");
        urlQueue.add("https://www.chinadaily.com.cn/business");
        urlQueue.add("https://www.cgtn.com/world");
        urlQueue.add("https://www.cgtn.com/business");

        // 中文种子URL
        urlQueue.add("https://news.sina.com.cn/world");
        urlQueue.add("https://news.163.com/world");
        urlQueue.add("https://world.people.com.cn");
    }

    private void crawlWorker() {
        try {
            while (true) {
                String url = urlQueue.poll(1, TimeUnit.SECONDS);
                if (url == null) {
                    if (englishDocCount.get() >= maxDocsPerLanguage &&
                            chineseDocCount.get() >= maxDocsPerLanguage) {
                        break;
                    }
                    continue;
                }

                if (!visitedUrls.add(url)) continue;

                String language = url.contains("en.people.cn") ||
                        url.contains("chinadaily.com.cn") ||
                        url.contains("cgtn.com") ? "english" : "chinese";

                if ((language.equals("english") && englishDocCount.get() >= maxDocsPerLanguage) ||
                        (language.equals("chinese") && chineseDocCount.get() >= maxDocsPerLanguage)) {
                    continue;
                }

                crawlPage(url, language);
                Thread.sleep(500 + (long)(Math.random() * 1000));
            }
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }
    }

    private void crawlPage(String url, String language) {
        try {
            Document doc = Jsoup.connect(url) .userAgent("Mozilla/5.0...").timeout(10000).get();
            // 保存原始文档
            String rawContent = doc.outerHtml();
            String savedPath = FileStorageUtil.saveHtmlDocument(
                    rawContent,
                    BASE_DATA_DIR,
                    language,
                    RAW_DIR,
                    "doc"
            );
            // 更新计数器
            if (language.equals("english"))   englishDocCount.incrementAndGet();
            else chineseDocCount.incrementAndGet();
            // 提取新链接
            if ((language.equals("english") && englishDocCount.get() < maxDocsPerLanguage) ||
                    (language.equals("chinese") && chineseDocCount.get() < maxDocsPerLanguage)) {

                doc.select("a[href]").stream()
                        .map(link -> link.attr("abs:href"))
                        .filter(this::isValidUrl)
                        .filter(newUrl -> !visitedUrls.contains(newUrl))
                        .limit(20) // 每页最多提取20个新链接
                        .forEach(urlQueue::offer);
            }
        } catch (IOException e) {
            System.err.println("Error crawling " + url + ": " + e.getMessage());
        }
    }

    private boolean isValidUrl(String url) {
        return url != null &&
                (url.startsWith("http://en.people.cn") ||
                        url.startsWith("https://www.chinadaily.com.cn") ||
                        url.startsWith("https://www.cgtn.com") ||
                        url.startsWith("https://news.sina.com.cn") ||
                        url.startsWith("https://news.163.com") ||
                        url.startsWith("https://world.people.com.cn"));
    }
}